import os
import numpy as np
from common.python.common import *


class ImageDecoder:
    def __init__(self):
        self.dec_img_info = None
        self.image_path = ""
        self.use_hw_dec = False

    def init(self, image_path, enable_hw_dec=True):
        self.image_path = image_path
        para = sdk.lyn_image_dec_para_t()

        para.img_url = image_path
        para.scale = sdk.lyn_scale_t.SCALE_NONE
        para.output_fmt = sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12
        para.align = True

        # sdk python中 没有SCALE_DOWN_8X
        vec_scale = [
            sdk.lyn_scale_t.SCALE_NONE,
            sdk.lyn_scale_t.SCALE_DOWN_2X,
            sdk.lyn_scale_t.SCALE_DOWN_4X,
        ]
        hard_dec_support = False

        for scale in vec_scale:
            para.scale = scale
            self.dec_img_info, hard_dec_support, ret = sdk.lyn_image_get_info(para)
            error_check(ret)

            if (self.dec_img_info.output.height <= 1080) and (
                self.dec_img_info.output.width <= 1920
            ):
                break

        if enable_hw_dec and hard_dec_support:
            self.use_hw_dec = True

    def get_image_dec_info(self):
        return self.dec_img_info

    def decode_image_to_device(self, stream, device_out_mem):
        with open(self.image_path, "rb") as image_file:
            image_read_np = np.frombuffer(image_file.read(), dtype=np.uint8)
        image_data_np = image_read_np.copy()
        image_size = len(image_data_np)
        image_data = sdk.lyn_numpy_to_ptr(image_data_np)

        dec_attr = sdk.lyn_image_dec_attr_t()
        dec_attr.in_buf.size = image_size
        dec_attr.info = self.dec_img_info
        dec_attr.out_buf.size = self.dec_img_info.output.predict_buf_size
        if self.use_hw_dec:
            dec_attr.in_buf.data, ret = sdk.lyn_malloc(dec_attr.in_buf.size)
            error_check(ret)
            dec_attr.out_buf.data = device_out_mem
            error_check(
                sdk.lyn_memcpy(
                    dec_attr.in_buf.data,
                    image_data,
                    dec_attr.in_buf.size,
                    sdk.lyn_memcpy_dir_t.ClientToServer,
                )
            )

            error_check(sdk.lyn_jpeg_decode_async(stream, dec_attr))
            error_check(sdk.lyn_synchronize_stream(stream))
            error_check(sdk.lyn_free(dec_attr.in_buf.data))
        else:
            dec_attr.in_buf.data = image_data
            dec_attr.out_buf.data = sdk.c_malloc(dec_attr.out_buf.size)
            error_check(sdk.lyn_image_decode_soft(dec_attr))
            error_check(
                sdk.lyn_memcpy(
                    device_out_mem,
                    dec_attr.out_buf.data,
                    dec_attr.out_buf.size,
                    sdk.lyn_memcpy_dir_t.ClientToServer,
                )
            )
            sdk.c_free(dec_attr.out_buf.data)
        return True
